import traceback

import glob
import os
import traceback

import cv2
import numpy as np
from PIL import Image


def webp_to_mp4(input_path, output_path, fps=24):
    """
    增强版：支持WebP/GIF混合转换

    参数更新：
    input_path: 可以是包含图片的目录，或单个GIF文件路径
    """
    # 新增GIF处理逻辑
    if os.path.isfile(input_path) and input_path.lower().endswith('.gif'):
        # 处理单个GIF文件的多帧
        gif = Image.open(input_path)
        width, height = gif.size
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        video = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
        frame_count = 0
        while True:
            try:
                gif.seek(frame_count)
                frame = cv2.cvtColor(np.array(gif.convert('RGB')), cv2.COLOR_RGB2BGR)
                video.write(frame)
                frame_count += 1
            except EOFError:
                print(f"EOFError occurred in single GIF processing:")
                print(traceback.format_exc())
                break
        video.release()
    else:
        # 原处理逻辑增加.gif支持
        webp_files = sorted(glob.glob(os.path.join(input_path, "*.[wW][eE][bB][pP]"))) + \
                     sorted(glob.glob(os.path.join(input_path, "*.[gG][iI][fF]")))

        # 新增：添加文件类型过滤和排序
        webp_files = [f for f in webp_files if f.lower().endswith(('.webp', '.gif'))]

        # 修改点：当输入的是目录时，自动合并目录下所有文件
        with Image.open(webp_files[0]) as img:
            width, height = img.size

        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        video = cv2.VideoWriter(output_path, fourcc, fps, (width, height))

        # 修改点：添加GIF多帧处理逻辑
        for file in webp_files:
            if file.lower().endswith('.gif'):
                gif = Image.open(file)
                gif_frames = 0
                while True:
                    try:
                        gif.seek(gif_frames)
                        frame = cv2.cvtColor(np.array(gif.convert('RGB')), cv2.COLOR_RGB2BGR)
                        video.write(frame)
                        gif_frames += 1
                    except EOFError:
                        print(f"EOFError occurred in {file}:")
                        print(traceback.format_exc())
                        break
            else:  # 处理WebP文件
                img = Image.open(file).convert('RGB')
                video.write(cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR))

        video.release()


if __name__ == '__main__':
    webp_to_mp4(
        input_path="C:/Users/Administrator/Desktop/测试数据集",
        output_path="output.mp4",
        fps=30
    )
